import cv2
#import fun as f
import sesson2.subject1.fun as f
import os
import numpy as np
import  time
import math
time_start=time.time()

# from math import *



def get_pq(a1, a2):
    ag1=np.arctan(a1[0]/a1[1])*180/np.pi
    ag2=np.arctan(a2[0]/a2[1])*180/np.pi
    return ag1-ag2

def calangle( kp1, kp2):
    anl=[]
    xa=[]
    xb=[]
    bll=[]
    bl=0
    da=0
    xiao=0
    for d in range(1, len(kp1)):
        jd = get_pq(kp1[d-1] - kp1[d], kp2[d-1] - kp2[d])
        xa1 = np.sqrt(np.sum(np.power((kp1[d-1] - kp1[d]), 2)))
        xb1 = np.sqrt(np.sum(np.power((kp2[d-1] - kp2[d]), 2)))
        if xa1!=0 and xb1!=0:
            bl=xb1/xa1
            bll.append(bl)
        if jd>0:
            da+=1
        else:
            xiao+=1
        anl.append(np.abs(jd))
    print("比例：",bll)
    if len(bll)>1:
        bll.remove(max(bll))
    bl=np.sum(bll)/len(bll)
    print("缩放比例：" + str(bl))
    print(anl)
    if len(anl)>1:
        anl.remove(max(anl))
    if len(anl)>1:
        anl.remove(max(anl))
    print(anl)
    jda=np.sum(np.abs(anl))/len(anl)
    if da <xiao:
        jda=-jda
    print("旋转角度：", jda)
    return bl,jda

# 去除黑边的操作
crop_image = lambda img, x0, y0, w, h: img[x0:x0+w, y0:y0+h] # 定义裁切函数，后续裁切黑边使用
crop_image1 = lambda img, x0, y0, w, h: img[x0:w, y0:h] # 定义裁切函数，后续裁切黑边使用
def rotate_image(img, angle, crop):
    """
    angle: 旋转的角度
    crop: 是否需要进行裁剪，布尔向量
    """
    w, h = img.shape[:2]
    # 旋转角度的周期是360°
    angle %= 360
    # 计算仿射变换矩阵
    M_rotation = cv2.getRotationMatrix2D((w / 2, h / 2), angle, 1)
    # 得到旋转后的图像
    img_rotated = cv2.warpAffine(img, M_rotation, (w, h))
    # img_rotated=rotate_bound(img, angle)
    # w, h = img_rotated.shape[:2]
    # 如果需要去除黑边
    if crop:
        # 裁剪角度的等效周期是180°
        angle_crop = angle % 180
        if angle > 90:
            angle_crop = 180 - angle_crop
        # 转化角度为弧度
        theta = angle_crop * np.pi / 180
        # 计算高宽比
        hw_ratio = float(h) / float(w)
        # 计算裁剪边长系数的分子项
        tan_theta = np.tan(theta)
        numerator = np.cos(theta) + np.sin(theta) * np.tan(theta)

        # 计算分母中和高宽比相关的项
        r = hw_ratio if h > w else 1 / hw_ratio
        # 计算分母项
        denominator = r * tan_theta + 1
        # 最终的边长系数
        crop_mult = numerator / denominator

        # 得到裁剪区域
        w_crop = int(crop_mult * w)
        h_crop = int(crop_mult * h)
        x0 = int((w - w_crop) / 2)
        y0 = int((h - h_crop) / 2)

        img_rotated = crop_image(img_rotated, x0, y0, w_crop, h_crop)

    return img_rotated



def rotate_bound(image, angle):
    # grab the dimensions of the image and then determine the
    # center
    (h, w) = image.shape[:2]
    (cX, cY) = (w // 2, h // 2)

    # grab the rotation matrix (applying the negative of the
    # angle to rotate clockwise), then grab the sine and cosine
    # (i.e., the rotation components of the matrix)
    M = cv2.getRotationMatrix2D((cX, cY), angle, 1.0)
    cos = np.abs(M[0, 0])
    sin = np.abs(M[0, 1])

    # compute the new bounding dimensions of the image
    nW = int((h * sin) + (w * cos))
    nH = int((h * cos) + (w * sin))

    # adjust the rotation matrix to take into account translation
    M[0, 2] += (nW / 2) - cX
    M[1, 2] += (nH / 2) - cY


    ima=cv2.warpAffine(image, M, (nW, nH))
#,borderMode=cv2.BORDER_REPLICATE
    # cv2.imshow("2",img_rotated)
    # perform the actual rotation and return the image
    return ima


def getkey(x1):
    return x1[2]




def yuan(sift_kp_cat,results,r, ff,flag):

    for i in range(0, len(sift_kp_cat)):

        for rr in range(len(results)):
            if i== results[rr]:
                flag=1
                break
        if flag==1:
            flag=0
            continue
        good_match = []
        matches = flann.knnMatch(sift_kp_cat[i][1], sift_kp_cat[results[r]][1], 2)

        for m in matches:
            # 当最近距离跟次近距离的比值小于ratio值时，保留此匹配对
            if len(m) == 2 and m[0].distance < m[1].distance * ratio:
                # 存储两个点在featuresA, featuresB中的索引值
                good_match.append((m[0].trainIdx, m[0].queryIdx, m[0].distance))
        print("good_match num:", len(good_match),"id:", i)
        # 当筛选后的匹配对大于4时，计算视角变换矩阵
        if len(good_match) > 4:
            # 求缩放比例
            little = sorted(good_match, key=getkey)
            pts11 = []
            pts22 = []
            pts1 = np.float32([sift_kp_cat[i][0][j] for (_, j, _) in little])
            pts2 = np.float32([sift_kp_cat[results[r]][0][j] for (j, _, _) in little])
            # print(pts1)
            # print(pts1)
            pts11.append(pts1[0])
            pts22.append(pts2[0])
            for z in range(1, len(pts1)):
                if (pts1[z][0] != pts11[len(pts11)-1][0] or pts1[z][1] != pts11[len(pts11)-1][1]) and (pts2[z][0] != pts22[len(pts22)-1][0] or pts2[z][1] != pts22[len(pts22)-1][1])  :
                    if len(pts11)>1:
                        if pts1[z][0] == pts11[len(pts11)-2][0] and pts1[z][1] == pts11[len(pts11)-2][1] and pts2[z][0] == pts22[len(pts22)-2][0] and pts2[z][1] == pts22[len(pts22)-2][1] :
                            continue
                    if len(pts11) == 4:
                        break
                    pts11.append(pts1[z])
                    pts22.append(pts2[z])
            print("len(pts11)",len(pts11))

            bl,ANGLE=calangle(pts11, pts22)
            result=imgnames[i]

            if ff:
                if ANGLE>5 or ANGLE <-5:
                    #完整但有黑边
                    if ANGLE <-1:
                         ANGLE=-ANGLE
                    # if len(ang)>1:
                    #     if ang[0] < 0:
                    #         if ANGLE>0:
                    #             ANGLE=-ANGLE
                    #     else:
                    #         if ANGLE<0:
                    #             ANGLE=-ANGLE
                    print(results[r],ANGLE,len(ang))
                    result = rotate_bound(imgnames[i], ANGLE)

                    # result= get_img_rot_broa(imgnames[i], ANGLE, filled_color=-1)

                    imgnames[i] = result
                    rgb = cv2.split(imgnames[i])
                    imgrgbs[i] = rgb  # 三通道分离
                    difbr = rgb[2] - rgb[0]
                    detesiftsrc[i] = difbr
                    (kps, features) = sift.detectAndCompute(difbr, None)
                    kps = np.float32([kp.pt for kp in kps])
                    sift_kp_cat[i] = (kps, features)
            else:
                if bl >= 1.05:
                    a0 = imgnames[i]
                    result = cv2.resize(a0, None, fx=bl, fy=bl, interpolation=cv2.INTER_CUBIC)
                    # cv2.imshow("a2", result)
                    # cv2.waitKey(0)
                    imgnames[i] = result
                    rgb = cv2.split(imgnames[i])
                    imgrgbs[i]=rgb  # 三通道分离
                    difbr = rgb[2] - rgb[0]
                    detesiftsrc[i]=difbr
                    (kps, features) = sift.detectAndCompute(difbr, None)
                    kps = np.float32([kp.pt for kp in kps])
                    sift_kp_cat[i] = (kps, features)
                if bl < 0.95:
                    a0 = imgnames[i]
                    result = cv2.resize(a0, None, fx=bl, fy=bl, interpolation=cv2.INTER_AREA)
                    # cv2.imshow("a2", a2)
                    imgnames[i] = result
                    rgb = cv2.split(imgnames[i])
                    imgrgbs[i] = rgb  # 三通道分离
                    difbr = rgb[2] - rgb[0]
                    detesiftsrc[i]=difbr
                    (kps, features) = sift.detectAndCompute(difbr, None)
                    kps = np.float32([kp.pt for kp in kps])
                    sift_kp_cat[i] = (kps, features)
            # result = f.cutpic(imgnames[i])
            cv2.imwrite((folder + f.PIC_NAME + str(i + 1) + ".png"), result)
            results.append(i)

            print("**************与" + str(results[r]) + "匹配的图像有"+str(i)+"*******************")
            #print(H)
        else:
            wei.append((i))
            print("wei.append((i)):",i)
            print("not enough matches are found--%d" % (len(good_match)))
            # 如果匹配对小于4时，返回None
            continue


def zai(sift_kp_cat, wei, ff,rat,num,flag):
    pp=[]
    for r in range(len(wei)):  # 就是说 从0开始遍历，可以与0匹配的图对应的id加入到results里，第二次循环的时候让与results【1】匹配的图对应的id加入到results里，
        print("**************与" + str(wei[r]) + "匹配的图像*******************")
        # wei = []
        for i in range(0, len(sift_kp_cat)):
            # for rr in range(len(wei)):
            if i == wei[r]:
                continue
                # flag = 1
                # break
            # if flag == 1:
            #     flag = 0
            #     continue
            good_match = []
            # matches = matcher.knnMatch(sift_kp_cat[wei[r]][1], sift_kp_cat[i][1], 2)
            matches = flann.knnMatch(sift_kp_cat[wei[r]][1], sift_kp_cat[i][1], 2)
            for m in matches:
                # 当最近距离跟次近距离的比值小于ratio值时，保留此匹配对
                if len(m) == 2 and m[0].distance < m[1].distance * rat:
                    # 存储两个点在featuresA, featuresB中的索引值
                    good_match.append((m[0].trainIdx, m[0].queryIdx, m[0].distance))
            print("good_match num:", len(good_match), "id:", i)
            # 当筛选后的匹配对大于4时，计算视角变换矩阵
            if len(good_match) > num:
                # 获取匹配对的点坐标
                ptsA = np.float32([sift_kp_cat[wei[r]][0][j] for (_, j, _) in good_match])
                ptsB = np.float32([sift_kp_cat[i][0][j] for (j, _, _) in good_match])

                # 求缩放比例
                little = sorted(good_match, key=getkey)
                pts11 = []
                pts22 = []
                a = 0
                b = 0
                pts1 = np.float32([sift_kp_cat[wei[r]][0][j] for (_, j, _) in little])
                pts2 = np.float32([sift_kp_cat[i][0][j] for (j, _, _) in little])
                # print(pts1)
                # print(pts1)
                pts11.append(pts1[0])
                pts22.append(pts2[0])
                for z in range(1, len(pts1)):
                    if (pts1[z][0] != pts11[len(pts11) - 1][0] or pts1[z][1] != pts11[len(pts11) - 1][1]) and (
                            pts2[z][0] != pts22[len(pts22) - 1][0] or pts2[z][1] != pts22[len(pts22) - 1][1]):
                        if len(pts11) > 1:
                            if pts1[z][0] == pts11[len(pts11) - 2][0] and pts1[z][1] == pts11[len(pts11) - 2][1] and \
                                    pts2[z][0] == pts22[len(pts22) - 2][0] and pts2[z][1] == pts22[len(pts22) - 2][1]:
                                continue
                        if len(pts11) == 4:
                            break
                        pts11.append(pts1[z])
                        pts22.append(pts2[z])
                print(pts11)

                bl, ANGLE = calangle(pts11, pts22)
                result = imgnames[wei[r]]

                if ff:
                    if ANGLE > 3 or ANGLE < -3:
                        # 完整但有黑边
                        if ANGLE < -1:
                            ANGLE = -ANGLE
                        # print(results[r], ANGLE, len(ang))
                        result = rotate_bound(imgnames[wei[r]], ANGLE)


                        imgnames[wei[r]] = result
                        rgb = cv2.split(imgnames[wei[r]])
                        imgrgbs[wei[r]] = rgb  # 三通道分离
                        difbr = rgb[2] - rgb[0]
                        detesiftsrc[wei[r]] = difbr
                        (kps, features) = sift.detectAndCompute(difbr, None)
                        kps = np.float32([kp.pt for kp in kps])
                        sift_kp_cat[wei[r]] = (kps, features)
                else:
                    if bl >= 1.05:
                        a0 = imgnames[wei[r]]
                        result = cv2.resize(a0, None, fx=bl, fy=bl, interpolation=cv2.INTER_CUBIC)
                        imgnames[wei[r]] = result
                        rgb = cv2.split(imgnames[wei[r]])
                        imgrgbs[wei[r]] = rgb  # 三通道分离
                        difbr = rgb[2] - rgb[0]
                        detesiftsrc[wei[r]] = difbr
                        (kps, features) = sift.detectAndCompute(difbr, None)
                        kps = np.float32([kp.pt for kp in kps])
                        sift_kp_cat[wei[r]] = (kps, features)
                    if bl < 0.95:
                        a0 = imgnames[wei[r]]
                        result = cv2.resize(a0, None, fx=bl, fy=bl, interpolation=cv2.INTER_AREA)
                        # cv2.imshow("a2", a2)
                        imgnames[wei[r]] = result
                        rgb = cv2.split(imgnames[i])
                        imgrgbs[wei[r]] = rgb  # 三通道分离
                        difbr = rgb[2] - rgb[0]
                        detesiftsrc[wei[r]] = difbr
                        (kps, features) = sift.detectAndCompute(difbr, None)
                        kps = np.float32([kp.pt for kp in kps])
                        sift_kp_cat[wei[r]] = (kps, features)
                cv2.imwrite((folder + f.PIC_NAME + str(wei[r] + 1) + ".png"), result)
                pp.append(wei[r])
                break
            else:
                print("not enough matches are found--%d" % (len(good_match)))
                # 如果匹配对小于4时，返回None
                continue
    return pp

imgnames = []
results = []
results1 = []
results2 = []
hui=[]
imgrgbs = []
detesiftsrc = []
reprojThresh = 4
# orb = cv2.ORB_create()
sift = cv2.xfeatures2d.SIFT_create()
# sift = cv2.xfeatures2d.SURF_create()
sift_kp_cat = []
# 设置FLANN匹配器参数
# algorithm设置可参考https://docs.opencv.org/3.1.0/dc/d8c/namespacecvflann.html
indexParams = dict(algorithm=0, trees=5)
searchParams = dict(checks=50)
# 定义FLANN匹配器
flann = cv2.FlannBasedMatcher(indexParams,searchParams)
id,pa=f.read_client_path()
folder = os.getcwd()[:-4] + 'new_folder/test/'

wei=[8]

ratio=0.4
ANGLE=0
matchindex = 0
matchnum = []
good_matchs = []
ang = []
# 获取此py文件路径，在此路径选创建在new_folder文件夹中的test文件夹

if not os.path.exists(folder):
    os.makedirs(folder)






def session2_1_main():

    # bf = cv2.BFMatcher_create(cv2.NORM_HAMMING, crossCheck=True)
    global wei
    for i in range(1, 10):

        filepath=(pa+f.PIC_NAME+str(i)+".png")
        # filepath = (f.PATH+f.PIC_NAME+str(i)+".png")
        imgrgb = cv2.imread(filepath)
        if i==1:
            cv2.imwrite((folder+ f.PIC_NAME + str(i) + ".png"), imgrgb)
            results.append(i-1)
            results1.append(i - 1)
            results2.append(i - 1)
        # imgrgb= np.array(imgrgb, dtype='uint8')
        imgnames.append(imgrgb)
        rgb = cv2.split(imgrgb)
        imgrgbs.append(rgb)  # 三通道分离
        difbr = rgb[2]-rgb[0]
        detesiftsrc.append(difbr)
        (kps, features) = sift.detectAndCompute(difbr, None)
        kps = np.float32([kp.pt for kp in kps])
        sift_kp_cat.append((kps,features))



    # wei=[8,9]
    print("!!!!!!!!!!!!!!!!!第一次匹配----旋转！！！！！！！！！！")
    for r in range(9):

        if r>len(results1)-1:
            print("没有匹配全")
            print("未匹配的有：", wei)
            break
        print("**************与"+str(results1[r])+"匹配的图像*******************")
        wei = []
        yuan(sift_kp_cat, results1,r, True,0)


    print("!!!!!!!!!!!!!!!!!第一次匹配----缩放！！！！！！！！！！")
    for r in range(9):

        if r > len(results2) - 1:
            print("没有匹配全")
            print("未匹配的有：", wei)
            break
        print("**************与" + str(results2[r]) + "匹配的图像*******************")
        wei = []
        yuan(sift_kp_cat, results2, r, False,0)

    print("!!!!!!!!!!!!!!!!!第二次匹配----旋转缩放！！！！！！！！！！")
    if len(wei)>0:
        zai(sift_kp_cat, wei,  True, 0.45, 4, 0)
        pp=zai(sift_kp_cat, wei, False, 0.45, 4, 0)
        print("已匹配的有：",pp)
        if len(pp)>0:
            for p in range(0,len(pp)):
                wei.remove(pp[p])
    print("未匹配的有：", wei)
    print("!!!!!!!!!!!!!!!!!第三次匹配----旋转缩放！！！！！！！！！！")
    if len(wei)>0:
        zai(sift_kp_cat, wei, True, 0.5, 5, 0)
        pp = zai(sift_kp_cat, wei, False, 0.5, 2, 0)
        # if len(pp) > 0:
        #     wei.remove(pp[0])
        if len(pp) > 0:
            for p in range( len(pp) ):
                wei.remove(pp[p])
    print("未匹配的有：",wei)

    imgname = []
    for i in range(1, 10):
        filepath = (folder + f.PIC_NAME + str(i) + ".png")
        imgrgb = cv2.imread(filepath)
        imgname.append(imgrgb)

    s = cv2.createStitcher()

    ret, p1 = s.stitch(
        [imgname[0], imgname[1], imgname[2], imgname[3], imgname[4], imgname[5], imgname[6], imgname[7],
         imgname[8]], None)

    print(ret)
    if ret==0:
        # cv2.namedWindow('1', cv2.WINDOW_NORMAL)
        # cv2.imshow('1', p1)
        cv2.imwrite(pa+"result.png", p1)
        pose, out = f.find_xys(p1)
        sum=len(pose)
        pp = []
        for p in pose:
            cv2.circle(out, (p[0], p[1]), 5, (0, 0, 255), 2)
            cv2.circle(p1, (p[0], p[1]), 5, (0, 0, 255), 2)
            # if(len(pp)<120):
            pp.append(p)
            # print(img[p[1], p[0]])
        # print("图片数据", pose)
        # cv2.imwrite("result.png", p1)
        # cv2.imshow("img0out1", p1)
        # cv2.imshow('img0out122', out)  # 显示原始图像
        # cv2.waitKey(0)
        # cv2.waitKey(0)
    else:
        sum=0
        for i in range(9):
            pose, out = f.find_xys(imgname[i])
            sum+=len(pose)
            pp = []
            for p in pose:
                cv2.circle(out, (p[0], p[1]), 5, (0, 0, 255), 2)
                cv2.circle(p1, (p[0], p[1]), 5, (0, 0, 255), 2)
                # if(len(pp)<120):
                pp.append(p)
            # cv2.imshow("img0out1",imgname[i])
            # cv2.imshow('img0out122', out)  # 显示原始图像
            # cv2.waitKey(0)
        sum=int(sum*0.64)
    f.toJons(sum,id,(pa+"result.png"))
    time_end=time.time()
    print('totally cost',time_end-time_start)
    return True


if __name__ == '__main__':
    session2_1_main()